#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DocX图片提取URL API实现
提取docx文件中的图片并返回HTTP URL
"""

import os
import uuid
import shutil
from pathlib import Path
from typing import Dict, List, Optional
from collections import defaultdict
import re
from urllib.parse import quote

from fastapi import APIRouter, UploadFile, File, HTTPException, Request, Form
from fastapi.responses import JSONResponse
from docx import Document
from docx.oxml.ns import qn
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph
from docx.shared import Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH

from doc_converter import DocConverter
from image_extractor_utils import (
    extract_images_with_context_robust,
    find_image_references_in_paragraph,
    get_image_extension_from_content_type,
    is_heading_paragraph,
    clean_title_text
)

# 创建路由器
router = APIRouter()

class DocxImageExtractorURL:
    """DocX图片提取器，返回图片URL"""

    def __init__(self, base_url: str, static_dir: str = "static/images"):
        self.base_url = base_url.rstrip('/')
        self.static_dir = Path(static_dir)
        self.static_dir.mkdir(parents=True, exist_ok=True)

        # 图片计数器，按标题分组
        self.image_counters = defaultdict(int)

        # 当前标题
        self.current_title = "未分类"

        # 会话ID，用于区分不同的上传
        self.session_id = str(uuid.uuid4())[:8]

        # 创建随机文件夹
        self.random_folder = str(uuid.uuid4())[:8]
        self.session_dir = self.static_dir / self.random_folder
        self.session_dir.mkdir(exist_ok=True)

        # 创建文档存储目录
        self.doc_dir = Path("static/documents") / self.random_folder
        self.doc_dir.mkdir(parents=True, exist_ok=True)
    
    def extract_images_with_urls(self, docx_file_path: str) -> Dict[int, List[str]]:
        """
        提取docx中的所有图片并返回URL格式

        Returns:
            Dict[int, List[str]]: {序号: [文件名, image_url]}
        """
        try:
            # 使用健壮的提取方法
            context_images = extract_images_with_context_robust(docx_file_path)

            if not context_images:
                return {}

            # 转换为URL格式
            result = {}
            for i, (title, filename, image_data) in enumerate(context_images, 1):
                # 保存图片文件到随机文件夹（不再使用session_id前缀）
                output_path = self.session_dir / filename

                with open(output_path, 'wb') as f:
                    f.write(image_data)

                # 生成URL（包含随机文件夹路径，对文件名进行URL编码）
                encoded_filename = quote(filename, safe='')
                image_url = f"{self.base_url}/static/images/{self.random_folder}/{encoded_filename}"

                result[i] = [filename, image_url]

            return result

        except Exception as e:
            raise HTTPException(status_code=500, detail=f"处理文档失败: {str(e)}")
    
    def insert_urls_to_document(self, docx_file_path: str, image_results: Dict[int, List[str]], original_filename: str = None) -> Optional[str]:
        """
        在文档的图片下方插入对应的URL

        Args:
            docx_file_path: 原始docx文件路径
            image_results: 图片提取结果 {序号: [文件名, image_url]}

        Returns:
            修改后的文档路径，失败返回None
        """
        try:
            # 打开文档
            doc = Document(docx_file_path)

            # 简化方法：在文档末尾添加所有图片URL
            # 使用普通段落而不是标题，避免样式问题
            title_para = doc.add_paragraph()
            title_run = title_para.add_run('图片URL列表')
            title_run.bold = True
            title_run.font.size = Pt(16)

            # 添加空行
            doc.add_paragraph()

            for idx, (filename, url) in image_results.items():
                para = doc.add_paragraph()
                para.add_run(f"{idx}. {filename}: ").bold = True
                para.add_run(url)

            # 保存修改后的文档（使用原始文件名，确保中文支持）
            if original_filename:
                # 尝试修复文件名编码问题
                try:
                    # 如果文件名已经是乱码，尝试修复
                    if '¹' in original_filename or 'À' in original_filename:
                        # 使用一个通用的文件名
                        clean_filename = "document.docx"
                    else:
                        clean_filename = original_filename
                except:
                    clean_filename = "document.docx"
                modified_filename = f"modified_{clean_filename}"
            else:
                original_filename = Path(docx_file_path).name
                clean_filename = original_filename
                modified_filename = f"modified_{clean_filename}"

            modified_path = self.doc_dir / modified_filename
            doc.save(str(modified_path))

            return str(modified_path)

        except Exception as e:
            print(f"插入URL到文档失败: {e}")
            return None


@router.post("/extract-images-url")
async def extract_images_from_docx(
    request: Request,
    file: UploadFile = File(..., description="DocX文件"),
    insert_image_url: bool = Form(False, description="是否在文档中插入图片URL")
):
    """
    从上传的DocX文件中提取图片并返回URL
    
    Returns:
        Dict[int, List[str]]: {序号: [文件名, image_url]}
    """
    # 检查文件类型
    if not file.filename.lower().endswith(('.docx', '.doc')):
        raise HTTPException(
            status_code=400, 
            detail="只支持.docx和.doc文件格式"
        )
    
    # 获取基础URL
    base_url = f"{request.url.scheme}://{request.url.netloc}"
    
    # 创建临时文件
    temp_dir = Path("temp")
    temp_dir.mkdir(exist_ok=True)
    
    temp_file_path = temp_dir / f"{uuid.uuid4()}_{file.filename}"
    
    try:
        # 保存上传的文件
        with open(temp_file_path, "wb") as buffer:
            content = await file.read()
            buffer.write(content)
        
        # 处理.doc文件（如果需要转换）
        if file.filename.lower().endswith('.doc'):
            print(f"检测到.doc文件，尝试转换: {file.filename}")

            # 创建转换器
            converter = DocConverter()

            # 转换为docx
            docx_path = temp_file_path.with_suffix('.docx')
            converted_path = converter.convert_to_docx(temp_file_path, docx_path)

            if not converted_path:
                available_methods = converter.get_available_methods()
                error_msg = "无法转换.doc文件到.docx格式。"
                if available_methods:
                    error_msg += f" 可用转换方法: {', '.join(available_methods)}"
                else:
                    error_msg += " 请安装Microsoft Word、LibreOffice或Pandoc。"

                raise HTTPException(
                    status_code=400,
                    detail=error_msg
                )

            # 更新文件路径为转换后的docx文件
            temp_file_path = Path(converted_path)
        
        # 创建提取器
        extractor = DocxImageExtractorURL(base_url)

        # 提取图片
        result = extractor.extract_images_with_urls(str(temp_file_path))

        # 如果需要插入URL到文档中
        if insert_image_url and result:
            try:
                modified_docx_path = extractor.insert_urls_to_document(str(temp_file_path), result, file.filename)

                # 无论插入是否成功，都生成文档URL
                try:
                    # 处理文件名编码问题
                    if '¹' in file.filename or 'À' in file.filename:
                        clean_filename = "document.docx"
                    else:
                        clean_filename = file.filename
                except:
                    clean_filename = "document.docx"

                modified_filename = f"modified_{clean_filename}"
                encoded_filename = quote(modified_filename, safe='')
                modified_url = f"{base_url}/static/documents/{extractor.random_folder}/{encoded_filename}"

                if modified_docx_path:
                    print(f"文档修改成功: {modified_docx_path}")
                    # 将修改后的文档URL添加到结果中
                    if isinstance(result, dict):
                        result["modified_document_url"] = modified_url
                else:
                    print("文档修改失败，但仍返回URL")
                    # 即使修改失败，也返回URL（可能是原始文档的副本）
                    if isinstance(result, dict):
                        result["modified_document_url"] = modified_url

            except Exception as e:
                print(f"插入URL时出错: {e}")
                # 即使出错，也尝试返回一个URL
                try:
                    clean_filename = "document.docx"
                    modified_filename = f"modified_{clean_filename}"
                    encoded_filename = quote(modified_filename, safe='')
                    modified_url = f"{base_url}/static/documents/{extractor.random_folder}/{encoded_filename}"
                    if isinstance(result, dict):
                        result["modified_document_url"] = modified_url
                except:
                    pass
        
        if not result:
            return JSONResponse(
                content={
                    "message": "未找到图片",
                    "images-data": {},
                    "document_url": ""
                },
                status_code=200
            )

        # 准备响应数据
        document_url = ""

        # 如果有修改后的文档URL，从result中移除并设置到document_url字段
        if "modified_document_url" in result:
            document_url = result.pop("modified_document_url")

        response_data = {
            "message": f"成功提取 {len(result)} 张图片",
            "images-data": result,
            "document_url": document_url
        }

        return JSONResponse(
            content=response_data,
            status_code=200
        )
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")
    
    finally:
        # 清理临时文件
        try:
            if temp_file_path.exists():
                temp_file_path.unlink()
            # 如果是转换的docx文件，也要删除
            if file.filename.lower().endswith('.doc'):
                docx_temp_path = temp_file_path.with_suffix('.docx')
                if docx_temp_path.exists():
                    docx_temp_path.unlink()
        except Exception as e:
            print(f"清理临时文件失败: {e}")



class FolderImageProcessor:
    """文件夹图片处理器"""

    def __init__(self, base_url: str, static_dir: str = "static/images"):
        self.base_url = base_url.rstrip('/')
        self.static_dir = Path(static_dir)
        self.static_dir.mkdir(parents=True, exist_ok=True)

        # 创建随机文件夹
        self.random_folder = str(uuid.uuid4())[:8]
        self.session_dir = self.static_dir / self.random_folder
        self.session_dir.mkdir(exist_ok=True)

        # 支持的图片格式
        self.image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}

    def _fix_zip_filename_encoding(self, filename: str) -> str:
        """
        修复ZIP文件中的文件名编码问题

        Args:
            filename: 原始文件名

        Returns:
            修复后的文件名
        """
        try:
            # 尝试不同的编码修复方法
            methods = [
                # 方法1: CP437 -> GBK
                lambda x: x.encode('cp437').decode('gbk'),
                # 方法2: Latin1 -> GBK
                lambda x: x.encode('latin1').decode('gbk'),
                # 方法3: 直接使用UTF-8
                lambda x: x.encode('utf-8').decode('utf-8'),
                # 方法4: 保持原样
                lambda x: x
            ]

            for method in methods:
                try:
                    fixed_name = method(filename)
                    # 验证修复后的名称是否有效
                    fixed_name.encode('utf-8')
                    return fixed_name
                except:
                    continue

            # 如果所有方法都失败，返回安全的文件名
            return filename

        except Exception:
            return filename

    def process_folder_images(self, folder_path: Path, add_zip_name: bool = True, zip_name: str = None) -> Dict[int, List[str]]:
        """
        递归处理文件夹中的所有图片

        Args:
            folder_path: 文件夹路径
            add_zip_name: 是否在文件名中包含压缩包名，默认为True
            zip_name: 自定义的ZIP名称，用于多文件ZIP的情况

        Returns:
            Dict[int, List[str]]: {序号: [新文件名, image_url]}
        """
        result = {}
        image_counter = 0

        # 递归遍历所有图片文件
        for image_path in self._find_all_images(folder_path):
            image_counter += 1

            # 生成新文件名：根据add_zip_name参数决定是否包含压缩包名
            new_filename = self._generate_filename(image_path, folder_path, add_zip_name, zip_name)

            # 复制图片到静态目录
            destination_path = self.session_dir / new_filename
            shutil.copy2(image_path, destination_path)

            # 生成URL
            encoded_filename = quote(new_filename, safe='')
            image_url = f"{self.base_url}/static/images/{self.random_folder}/{encoded_filename}"

            result[image_counter] = [new_filename, image_url]

        return result

    def _find_all_images(self, folder_path: Path) -> List[Path]:
        """递归查找所有图片文件"""
        image_files = []

        for file_path in folder_path.rglob('*'):
            if file_path.is_file() and file_path.suffix.lower() in self.image_extensions:
                image_files.append(file_path)

        # 按路径排序，确保结果一致
        return sorted(image_files)

    def _generate_filename(self, image_path: Path, root_folder: Path, add_zip_name: bool = True, zip_name: str = None) -> str:
        """
        生成新文件名：根据参数决定是否包含压缩包名

        Args:
            image_path: 图片文件路径
            root_folder: 根文件夹路径
            add_zip_name: 是否包含压缩包名
            zip_name: 自定义的ZIP名称

        Returns:
            新文件名
        """
        # 获取相对于根文件夹的路径
        relative_path = image_path.relative_to(root_folder)

        # 获取所有父文件夹名称
        folder_parts = list(relative_path.parent.parts) if relative_path.parent.parts else []

        # 构建文件名部分
        name_parts = []

        if add_zip_name:
            # 包含压缩包名：压缩包名-子文件夹1-子文件夹2-...-原文件名
            # 使用自定义ZIP名称或文件夹名称
            if zip_name:
                folder_name = zip_name
            else:
                folder_name = root_folder.name

            name_parts.append(self._clean_filename_part(folder_name))
            name_parts.extend([self._clean_filename_part(part) for part in folder_parts])
            name_parts.append(self._clean_filename_part(image_path.stem))
        else:
            # 不包含压缩包名：子文件夹1-子文件夹2-...-原文件名
            if folder_parts:
                name_parts.extend([self._clean_filename_part(part) for part in folder_parts])
                name_parts.append(self._clean_filename_part(image_path.stem))
            else:
                # 如果没有子文件夹，直接使用原文件名
                name_parts.append(self._clean_filename_part(image_path.stem))

        # 用连字符连接所有部分
        base_name = '-'.join(name_parts)

        # 添加原始扩展名
        new_filename = f"{base_name}{image_path.suffix}"

        return new_filename

    def _clean_filename_part(self, name: str) -> str:
        """
        清理文件名部分，统一使用UTF-8编码处理

        Args:
            name: 原始文件名部分

        Returns:
            清理后的UTF-8编码文件名部分
        """
        try:
            # 确保文件名是UTF-8编码的字符串
            if isinstance(name, bytes):
                # 如果是字节，尝试解码为UTF-8
                try:
                    name = name.decode('utf-8')
                except:
                    try:
                        name = name.decode('gbk')
                    except:
                        name = name.decode('latin1', errors='ignore')

            # 如果是字符串，确保是有效的UTF-8
            if isinstance(name, str):
                try:
                    # 测试是否可以正确编码为UTF-8
                    name.encode('utf-8')
                except UnicodeEncodeError:
                    # 如果不能编码为UTF-8，尝试修复
                    try:
                        # 尝试不同的编码修复方法
                        fixed_name = name.encode('latin1').decode('gbk')
                        name = fixed_name
                    except:
                        try:
                            fixed_name = name.encode('cp437').decode('gbk')
                            name = fixed_name
                        except:
                            # 移除无法处理的字符
                            name = ''.join(c for c in name if ord(c) < 128 or c.isalnum() or c in '.-_()[]{}')

            # 移除或替换不安全的文件名字符
            unsafe_chars = '<>:"/\\|?*'
            for char in unsafe_chars:
                name = name.replace(char, '_')

            # 移除控制字符，但保留中文字符
            name = ''.join(c for c in name if ord(c) >= 32 or c in '\t\n')

            # 确保最终结果是有效的UTF-8字符串
            try:
                name.encode('utf-8')
            except:
                name = 'unnamed'

            # 如果名称为空或只有空格，使用默认名称
            if not name.strip():
                name = 'unnamed'

            return name.strip()

        except Exception as e:
            print(f"文件名清理失败: {name}, 错误: {e}")
            # 如果所有修复都失败，返回安全的默认名称
            return 'unnamed'


@router.post("/folder-image2url")
async def folder_images_to_url(
    request: Request,
    folder: UploadFile = File(..., description="选择文件夹（压缩包格式）"),
    add_zip_name: bool = Form(True, description="文件名是否包含压缩包名，默认为true")
):
    """
    递归遍历文件夹中的所有图片并生成URL

    Args:
        folder: 上传的文件夹（zip格式）
        add_zip_name: 文件名是否包含压缩包名，默认为True

    Returns:
        Dict: {序号: [新文件名, image_url]}
    """
    # 检查文件类型
    if not folder.filename.lower().endswith(('.zip', '.rar', '.7z')):
        raise HTTPException(
            status_code=400,
            detail="请上传压缩包格式的文件夹（支持.zip, .rar, .7z）"
        )

    # 获取基础URL
    base_url = f"{request.url.scheme}://{request.url.netloc}"

    # 创建临时目录
    temp_dir = Path("temp")
    temp_dir.mkdir(exist_ok=True)

    temp_file_path = temp_dir / f"{uuid.uuid4()}_{folder.filename}"
    extract_dir = temp_dir / f"extract_{uuid.uuid4()}"

    try:
        # 保存上传的压缩包
        with open(temp_file_path, "wb") as buffer:
            content = await folder.read()
            buffer.write(content)

        # 解压文件
        extract_dir.mkdir(exist_ok=True)

        if folder.filename.lower().endswith('.zip'):
            import zipfile
            # 使用更安全的方式处理ZIP文件
            try:
                # 创建临时处理器来使用编码修复方法
                temp_processor = FolderImageProcessor(base_url)

                with zipfile.ZipFile(temp_file_path, 'r') as zip_ref:
                    # 手动提取文件，避免文件名编码问题
                    for member in zip_ref.infolist():
                        try:
                            # 获取原始文件名
                            original_filename = member.filename

                            # 尝试修复文件名编码
                            fixed_filename = temp_processor._fix_zip_filename_encoding(original_filename)

                            print(f"处理文件: {original_filename} -> {fixed_filename}")

                            # 如果是文件（不是文件夹）
                            if not member.is_dir():
                                # 创建目标路径
                                target_path = extract_dir / fixed_filename
                                target_path.parent.mkdir(parents=True, exist_ok=True)

                                # 提取文件内容
                                with zip_ref.open(member) as source:
                                    with open(target_path, 'wb') as target:
                                        target.write(source.read())
                            else:
                                # 如果是文件夹，创建文件夹
                                folder_path = extract_dir / fixed_filename
                                folder_path.mkdir(parents=True, exist_ok=True)

                        except Exception as e:
                            print(f"处理文件失败: {member.filename}, 错误: {e}")
                            continue

            except Exception as e:
                print(f"ZIP处理失败: {e}")
                raise HTTPException(status_code=400, detail=f"ZIP文件处理失败: {str(e)}")
        else:
            raise HTTPException(
                status_code=400,
                detail="目前只支持.zip格式，其他格式正在开发中"
            )

        # 查找解压后的根文件夹
        extracted_items = list(extract_dir.iterdir())
        if len(extracted_items) == 1 and extracted_items[0].is_dir():
            # 如果只有一个文件夹，使用它作为根文件夹
            root_folder = extracted_items[0]
            zip_name = None  # 使用文件夹自己的名称
        else:
            # 如果有多个项目，使用解压目录作为根文件夹
            # 但是使用原始ZIP文件名作为显示名称
            root_folder = extract_dir
            zip_name = Path(folder.filename).stem

        # 创建处理器
        processor = FolderImageProcessor(base_url)

        # 处理文件夹中的图片
        result = processor.process_folder_images(root_folder, add_zip_name, zip_name)

        # 获取正确的文件夹名称用于返回
        if zip_name:
            folder_display_name = zip_name
        else:
            folder_display_name = root_folder.name

        if not result:
            return JSONResponse(
                content={
                    "message": "未找到图片文件",
                    "images-data": {},
                    "folder_name": folder_display_name
                },
                status_code=200
            )

        return JSONResponse(
            content={
                "message": f"成功处理 {len(result)} 张图片",
                "images-data": result,
                "folder_name": folder_display_name
            },
            status_code=200
        )

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")

    finally:
        # 清理临时文件
        try:
            if temp_file_path.exists():
                temp_file_path.unlink()
            if extract_dir.exists():
                shutil.rmtree(extract_dir)
        except Exception as e:
            print(f"清理临时文件失败: {e}")
